Carry reward and per-token model_version, log reward/version metrics#562
Merged
Conversation
Add two optional per-token fields to the RL streaming schema and thread them through the data pipeline alongside `advantages` / `old_log_probabilities` (reusing `TokenDataDocument`/`TokenDataBatch`): - `reward`: the raw (un-normalized) reward, a per-rollout scalar broadcast per-token — distinct from the group-relative `advantage`. - `model_version`: the model version each token was generated under (documents-seen units), one per token, for staleness diagnostics. Both are optional (absent when the producer does not send them), so the batch/target plumbing guards on presence. The shared policy-gradient loss logs mean/max/min of each supplied field when `metrics != none` (GRPO and GSPO), registered only when the data is present. Staleness is `documents_seen - model_version`, derivable from the logged version stats and `documents_seen`. Co-Authored-By: Claude Opus 4.8 <noreply@anthropic.com>
…rim comments Inline the one-caller `_get_optional_target` into `_register_data_metrics`, dropping its dead `targets is None` disjunct (`to_kwargs` always sets the key); build the max/min sentinels once in the first-iteration init block; narrow the `losses` annotation to `dict`; move `_DATA_METRIC_FIELDS` above `__init__`; drop two restating/downstream-referencing comments. Co-Authored-By: Claude Opus 4.8 <noreply@anthropic.com>
The reward is averaged over the sample-filtered training batch, so its mean is biased and not a valid policy-performance metric. Rename only the logged metric label (the data field and kwargs key stay `reward`), and note the caveat. Co-Authored-By: Claude Opus 4.8 <noreply@anthropic.com>
…ersion Staleness must be computed per-step (model-version metrics are summed across steps and averaged only at log time, so it can't be derived from a single documents_seen afterwards). Thread the step's documents_seen into the loss kwargs through the same channel iteration uses (run_step -> BatchContext -> extra_kwargs), and subtract it from the per-token model version so the loss logs staleness directly. Other run_step callers take the default 0. Co-Authored-By: Claude Opus 4.8 <noreply@anthropic.com>
Match the file's established name for mask/labeled-token-count weight, and align the loop variable with the tuple's `data_key` comment. Co-Authored-By: Claude Opus 4.8 <noreply@anthropic.com>
Registering them conditionally on per-batch data presence made the metric's cross-DP all-reduce data-dependent: with metrics on (pipeline_parallel forced to 1), a rank that registered the metric all-reduces while a rank that skipped it does not, so mixed presence across DP ranks would hang. reward/model_version are always sent when these metrics are enabled, so register unconditionally like the other policy metrics; absence now fails uniformly instead of hanging. Hoist the shared per-document weight and extrema sentinels out of the loop. Co-Authored-By: Claude Opus 4.8 <noreply@anthropic.com>
Compute staleness (documents_seen - model_version) in the integer domain before casting to float, so the small difference isn't lost to float32 rounding when both counts exceed ~2^24. Also spell out neg/pos infinity sentinels, and drop a code-restating comment and a downstream-consumer reference. Co-Authored-By: Claude Opus 4.8 <noreply@anthropic.com>
This was referenced Jul 10, 2026
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
1 participant
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
Claude Sonnet 5, on behalf of @jlamypoirier.
Summary
Extracted from #553, as an independent piece with no dependency on the
documents_seen/weights_readyPRs (#559, #561) — verified by cherry-picking this commit directly ontomainwith no involvement of the other two. The "staleness" framing (
documents_seen - model_version)is a usage note for whoever reads the logged metrics; the code here doesn't compute or need
documents_seenitself.Was briefly combined with the
weights_readybroadcast in #557 (now closed) before splittinginto fully independent pieces.
pipeline alongside
advantages/old_log_probabilities(reusingTokenDataDocument/TokenDataBatch):reward: the raw (un-normalized) reward, a per-rollout scalar broadcast per-token — distinctfrom the group-relative
advantage.model_version: the model version each token was generated under (documents-seen units),one per token, for staleness diagnostics.
guards on presence.
metrics != none(GRPO and GSPO).
Paired PipelineRL change (raw-reward forwarding) is in a separate PR against PipelineRL's
fast-llmbranch.Tests
tests/data/test_streaming.py(schema round-trip) andtests/layers/test_lm_losses.py(metric registration, single-process and distributed): 584 passed / 21 skipped, 23 passed.
model_versionconsumer and its metric are in place, but the producer(tagging tokens with the active version inside vLLM's output path) is a separate,
cluster-validated follow-up; until it lands, the
model_version/staleness metrics simply don'tpopulate.